﻿using AdvVerifyGUI.Config;
using AdvVerifyGUI.Display;
using AdvVerifyGUI.Repositories;
using AdvVerifyGUI.Stage;
using AdvVerifyGUI.Utils.Image;
using AdvVerifyModels.Entities;
using AdvVerifyModels.Models;
using AdvVerifyModels.Utils;
using HWModule;
using Microsoft.EntityFrameworkCore;
using NLog;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Channels;
using System.Threading.Tasks;
using System.Windows;

namespace AdvVerifyGUI.Job;

public interface IJobQueue
{
    IJob Current { get; }

    IJob Pending { get; }

    bool IsBusy { get; }

    ISubscription<PatternResult> ResultUpdate { get; }

    IJob UnsafeRun(JobPlan plan, JobRunOptions options);

    Task<WaferNavigatorOptions> TestAlignment(IWaferNavigator nav, Pattern pattern, JobRunOptions options);

    Task<WaferNavigatorOptions> TestAlignment(WaferMap map, List<Pattern> patterns, JobRunOptions options);

    Task TestMeasurement(JobPlan plan, JobRunOptions options);

    // FIX ME
    Task MoveTo(JobPlan plan, int scanIndex)
    {
        return MoveTo(plan.ScanPoints[scanIndex]);
    }

    Task MoveTo(ScanPoint point);

    IAsyncEnumerable<Measurement> RetryJob(int jobId, JobPlan plan, List<int> scanIndices, JobRunOptions options);

    Task<Grid> MakeGrid(JobPlan plan, JobRunOptions options);
}

public class JobQueue(FrameDispatcher dispatcher,
    IStageService stageService,
    IImagingService imagingService,
    ISEMClient semClient,
    JobResultRepo repo) : IJobQueue
{
    static readonly Logger Logger = LogManager.GetCurrentClassLogger();

    readonly object lockObj = new();
    readonly MessageQueue<PatternResult> queue = new();
    public ISubscription<PatternResult> ResultUpdate => queue.Subscribe();

    bool isBusy;

    IJob current;
    JobContext pending;

    public IJob Current
    {
        get
        {
            lock (lockObj)
            {
                return current;
            }
        }
    }

    public IJob Pending
    {
        get
        {
            lock (lockObj)
            {
                return pending;
            }
        }
    }

    public bool IsBusy
    {
        get
        {
            lock (lockObj)
            {
                return isBusy;
            }
        }
    }

    public IJob UnsafeRun(JobPlan plan, JobRunOptions options)
    {
        var ctx = plan.Environment.RequireConfirmation ? NewPendingContext(options) : NewContext(options);
        _ = Execute(ctx, RunJobPlan(ctx, plan), Task.Run(() => RunNewJobProcessingPipeline(ctx, plan, options)));
        return ctx;
    }

    public async Task<WaferNavigatorOptions> TestAlignment(IWaferNavigator nav, Pattern pattern, JobRunOptions options)
    {
        var ctx = NewContext(options);
        return await Run(ctx, options, () => Align(ctx, nav, pattern))
            .ConfigureAwait(false);
    }

    public async Task<WaferNavigatorOptions> TestAlignment(WaferMap map, List<Pattern> patterns, JobRunOptions options)
    {
        var ctx = NewContext(options);
        return await Run(ctx, options, () => Align(ctx, map, patterns))
            .ConfigureAwait(false);
    }

    public async Task TestMeasurement(JobPlan plan, JobRunOptions options)
    {
        var ctx = NewContext(options);
        var imaging = Task.Run(() => RunImagingPipeline(ctx,
            stageService.Navigator,
            plan.ScanPoints,
            plan.Environment));
        var processing = Task.Run(() => RunNewJobProcessingPipeline(ctx, plan, options));
        await Execute(ctx, imaging, processing).ConfigureAwait(false);
    }

    public async Task MoveTo(ScanPoint p)
    {
        var options = new JobRunOptions
        {
            Id = -1,
            SkipMeasurement = true
        };
        var ctx = NewContext(options);
        await Run(ctx, options, async () =>
        {
            var pipeline = new ImagingPipeline(ctx, stageService.Navigator, semClient, imagingService)
            {
                SkipAutoFocus = true,
            };
            var c = new Coordinate(p.X, p.Y, p.Dx, p.Dy, p.AddrOffsetX, p.AddrOffsetY);
            await pipeline.MoveTo(c, p.Pattern).ConfigureAwait(false);
            return 0;
        });
    }

    // FIXME
    public IAsyncEnumerable<Measurement> RetryJob(int jobId, JobPlan plan, List<int> scanIndices, JobRunOptions options)
    {
        var ctx = NewContext(options);
        var chan = Channel.CreateUnbounded<Measurement>();

        // modify the plan in place
        var s = new List<ScanPoint>();
        foreach (var i in scanIndices)
        {
            s.Add(plan.ScanPoints[i]);
        }
        plan.Alignment = [];
        plan.ScanPoints = s;

        _ = Execute(ctx,
            RunJobPlan(ctx, plan),
            RunRetryProcessingPipeline(ctx, chan.Writer, jobId, plan, scanIndices, options));
        return chan.Reader.ReadAllAsync();
    }

    public async Task<Grid> MakeGrid(JobPlan plan, JobRunOptions options)
    {
        var ctx = NewContext(options);
        var env = plan.Environment;
        return await Run(ctx, options, async () =>
        {
            var navOptions = await Align(ctx, plan.WaferMap, plan.Alignment).ConfigureAwait(false);
            var nav = stageService.WaferNavigator(navOptions);
            var spec = nav.Options.Map;
            var w = spec.ChipW;
            var h = spec.ChipH;
            var grid = new Grid()
            {
                X = spec.ArrayX,
                Y = spec.ArrayY,
                Width = w,
                Height = h,
                Coordinates = new double[spec.ArrayX * spec.ArrayY * 2],
                Angle = nav.Options.Angle,
            };
            var pattern = plan.Alignment[1];
            var c = pattern.Coordinates[0];
            if (plan.ScanPoints.Count > 0)
            {
                var p = plan.ScanPoints[0];
                c = new Coordinate(p.X, p.Y, p.Dx, p.Dy, p.AddrOffsetX, p.AddrOffsetY);
                pattern = p.Pattern;
            }

            // assume we own the plan
            pattern.Kind = 2;
            pattern.Reference.DoCentering = true;

            var pipeline = ImagingPipeline(ctx, nav, env);
            for (var i = 0; i < spec.ArrayX; i++)
            {
                for (var j = 0; j < spec.ArrayY; j++)
                {
                    c = c with
                    {
                        X = i,
                        Y = j
                    };

                    if (ctx.IsCanceled)
                    {
                        throw new OperationCanceledException();
                    }

                    if (!spec.IsInside(c.X, c.Y, c.Dx, c.Dy) ||
                        !spec.IsInside(c.X, c.Y, c.Dx + c.AddrOffsetX, c.Dy + c.AddrOffsetY))
                    {
                        continue;
                    }

                    var mp = await pipeline.MoveTo(c, pattern);
                    var (sx, sy) = await nav.MotionControl.GetCoordinate();

                    if (mp.Score < mp.Pattern.Reference.Threshold)
                    {
                        continue;
                    }

                    var cond = semClient.Condition;
                    var mat = Transformation.NewTransformation(cond.PixelSize, cond.ImageAngle);
                    var offset = mat.Transform(new Vector(-mp.MatchDx, -mp.MatchDy));
                    grid[i, j] = new Vector(sx, sy) + offset;
                }
            }

            return grid;
        });
    }

    JobContext NewContext(JobRunOptions options)
    {
        lock (lockObj)
        {
            if (isBusy)
            {
                throw new InvalidOperationException("Job queue is busy.");
            }

            var ctx = new JobContext(this, dispatcher, options, false);
            current = ctx;
            isBusy = true;
            return ctx;
        }
    }

    JobContext NewPendingContext(JobRunOptions options)
    {
        lock (lockObj)
        {
            if (isBusy)
            {
                throw new InvalidOperationException("Job queue is busy.");
            }

            if (pending != null && !pending.Confirmation.IsCompleted)
            {
                throw new InvalidOperationException("Confirmation pending");
            }

            var ctx = new JobContext(this, dispatcher, options, true);
            current = ctx;
            pending = ctx;
            isBusy = true;
            return ctx;
        }
    }

    async Task<T> Run<T>(JobContext ctx, JobRunOptions options, Func<Task<T>> runJob)
    {
        return await Execute(ctx, runJob(), RunProcessingPipeline(ctx, options)).ConfigureAwait(false);
    }

    async Task<T> Execute<T>(JobContext ctx, Task<T> imaging, Task processing)
    {
        try
        {
            var task = await Task.WhenAny(imaging, processing).ConfigureAwait(false);
            if (task != imaging)
            {
                await task.ConfigureAwait(false);
            }

            var result = await imaging.ConfigureAwait(false);
            ctx.Complete();
            return result;
        }
        catch (Exception ex)
        {
            ctx.InternalCancel(ex);
            throw;
        }
        finally
        {
            lock (lockObj)
            {
                isBusy = false;
            }
        }
    }

    async Task RunProcessingPipeline(JobContext ctx, JobRunOptions options)
    {
        await foreach (var _ in RunPublisherPipeline(ctx, options).ConfigureAwait(false))
        {
            // ignored
        }
    }

    async IAsyncEnumerable<(int, ImageKind, string, double)> RunPublisherPipeline(JobContext ctx, JobRunOptions options)
    {
        using var pub = new Publisher(queue, options);
        await foreach (var (pImg, values) in ctx.RunProcessingPipeline())
        {
            switch (pImg.Kind)
            {
                case PatternKind.High:
                    {
                        var path = pub.PublishHigh(pImg);
                        yield return (pImg.Index, ImageKind.High, path, 0);
                    }
                    break;
                case PatternKind.Align:
                    {
                        var (raw, _) = values[0];
                        var path = pub.PublishAlign(pImg, raw);
                        yield return (pImg.Index, ImageKind.Align, path, 0);
                    }
                    break;
                case PatternKind.Addr:
                    {
                        var (raw, _) = values[0];
                        var path = pub.PublishAddr(pImg, raw);
                        yield return (pImg.Index, ImageKind.Addr, path, 0);
                    }
                    break;
                case PatternKind.Mp:
                    {
                        var path = pub.PublishMp(pImg, values);
                        yield return (pImg.Index, ImageKind.Raw, path, 0);
                        for (var i = 0; i < values.Length; i++)
                        {
                            var (img, v) = values[i];
                            path = pub.SaveResultImage(pImg.Pattern.Names[i], pImg.Coordinate, img);
                            yield return (pImg.Index, ImageKind.MP, path, v);
                        }
                    }
                    break;
            }
        }
    }

    async Task RunRetryProcessingPipeline(JobContext ctx,
        ChannelWriter<Measurement> writer,
        int jobId,
        JobPlan plan,
        List<int> scanIndices,
        JobRunOptions options)
    {
        try
        {
            using var dbCtx = repo.NewContext();
            var dict = await dbCtx.PatternCopies
                .Where(c => c.JobResultId == jobId)
                .Select(c => new
                {
                    c.Id,
                    c.Ordinal
                })
                .ToDictionaryAsync(static v => v.Ordinal, static v => v.Id)
                .ConfigureAwait(false);
            var i = -1;
            var index = -1;
            var rawIndex = 0;
            var mIndex = 0;
            await foreach (var (idx, kind, path, v) in RunPublisherPipeline(ctx, options).ConfigureAwait(false))
            {
                if (idx != index)
                {
                    i += 1;
                    rawIndex = 0;
                    mIndex = 0;
                    index = idx;

                    await dbCtx.Images
                        .Where(img => img.JobResultId == idx && img.ScanIndex == scanIndices[i])
                        .ExecuteDeleteAsync()
                        .ConfigureAwait(false);
                }

                var scanIndex = scanIndices[i];
                switch (kind)
                {
                    case ImageKind.Addr:
                        dbCtx.Images.Add(new PatternImage
                        {
                            Kind = ImageKind.Addr,
                            ScanIndex = scanIndex,
                            Filename = path,
                            JobResultId = jobId,
                        });
                        break;
                    case ImageKind.MP:
                        {
                            var pattern = plan.ScanPoints[i].Pattern;
                            var paramId = dict[pattern.Ordinal];
                            var m = await dbCtx.Measurements
                                .Where(m => m.ParamaterId == paramId && m.ScanIndex == scanIndex)
                                .AsTracking()
                                .FirstAsync()
                                .ConfigureAwait(false);
                            m.Value = v;
                            m.Status = MeasurementStatus.Finished;
                            dbCtx.Images.Add(new PatternImage
                            {
                                Kind = ImageKind.MP,
                                ScanIndex = scanIndex,
                                Index = mIndex,
                                Filename = path,
                                JobResultId = jobId,
                            });
                            writer.TryWrite(m);
                            mIndex += 1;
                        }
                        break;
                    case ImageKind.Raw:
                        {
                            var pattern = plan.ScanPoints[i].Pattern;
                            if (pattern.CDPoints.Count == 0)
                            {
                                var paramId = dict[pattern.Ordinal];
                                var m = await dbCtx.Measurements
                                   .Where(m => m.ParamaterId == paramId && m.ScanIndex == scanIndex)
                                   .AsTracking()
                                   .FirstAsync()
                                   .ConfigureAwait(false);
                                m.Status = MeasurementStatus.Finished;
                                writer.TryWrite(m);
                            }
                            dbCtx.Images.Add(new PatternImage
                            {
                                Kind = ImageKind.Raw,
                                ScanIndex = scanIndex,
                                Index = rawIndex,
                                Filename = path,
                                JobResultId = jobId,
                            });
                            rawIndex += 1;
                        }
                        break;
                    default:
                        throw new InvalidOperationException();
                }
            }

            await dbCtx.SaveChangesAsync().ConfigureAwait(false);

            writer.Complete();
        }
        catch (Exception ex)
        {
            writer.Complete(ex);
            throw;
        }
    }

    async Task RunNewJobProcessingPipeline(JobContext ctx, JobPlan plan, JobRunOptions options)
    {
        var patterns = new List<PatternCopy>();
        foreach (var p in plan.Alignment)
        {
            patterns.Add(new PatternCopy
            {
                Kind = p.Kind,
                Ordinal = p.Ordinal,
                Pattern = Copy(p)
            });
        }

        var scanOrder = 0;
        var dict = new Dictionary<int, PatternCopy>();
        var parameters = new Dictionary<(int, int), Paramater>();
        foreach (var s in plan.ScanPoints)
        {
            var pattern = s.Pattern;
            if (!dict.TryGetValue(s.PatternId, out var copy))
            {
                copy = new PatternCopy
                {
                    Kind = pattern.Kind,
                    Ordinal = pattern.Ordinal,
                    Pattern = Copy(s.Pattern)
                };
                dict[s.PatternId] = copy;

                if (pattern.CDPoints.Count == 0)
                {
                    parameters[(s.PatternId, 0)] = new Paramater
                    {
                        Name = $"P{s.Pattern.Ordinal}",
                        PatternCopy = copy,
                    };
                }
                else
                {
                    for (var i = 0; i < pattern.CDPoints.Count; i++)
                    {
                        parameters[(s.PatternId, i)] = new Paramater
                        {
                            Name = pattern.CDPoints[i].Name,
                            Index = i,
                            PatternCopy = copy,
                        };
                    }
                }
            }

            var n = pattern.CDPoints.Count == 0 ? 1 : pattern.CDPoints.Count;
            for (int i = 0; i < n; i++)
            {
                var param = parameters[(s.PatternId, i)];
                param.Measurements.Add(new Measurement
                {
                    ScanIndex = scanOrder,
                    X = s.X,
                    Y = s.Y,
                    Status = MeasurementStatus.NotRun,
                });
            }

            scanOrder += 1;
        }

        patterns.AddRange(dict.Values);

        using var writer = new JobResultWriter(repo.NewContext());
        var jobResult = new JobResult
        {
            Start = DateTimeOffset.Now.ToUnixTimeMilliseconds(),
            IsConfirmed = !plan.Environment.RequireConfirmation,
            Name = plan.Name,
            Operator = plan.Operator,
            LotId = plan.LotId,
            Process = plan.Process,
            WaferId = plan.WaferId,
            ConditionNo = plan.ConditionNo,
            RecipeId = plan.Id,
            OutputDirectory = options.OutputDirectory,
            Paramaters = [.. parameters.Values],
            Patterns = patterns,
        };
        await writer.Write(jobResult).ConfigureAwait(false);

        if (plan.Environment.RequireConfirmation)
        {
            // wait for user confirmation
            var id = jobResult.Id;
            _ = Task.Run(async () =>
            {
                await ctx.Confirmation.ConfigureAwait(false);
                var dbCtx = repo.NewContext();
                await dbCtx.JobResults
                    .Where(r => r.Id == id)
                    .ExecuteUpdateAsync(setters => setters
                        .SetProperty(r => r.IsConfirmed, true))
                    .ConfigureAwait(false);
            });
        }

        var high = 0;
        var align = 0;
        var index = -1;
        var scanIndex = -1;
        var rawIndex = 0;
        var mIndex = 0;
        try
        {
            await foreach (var (idx, kind, path, v) in RunPublisherPipeline(ctx, options).ConfigureAwait(false))
            {
                if (index != idx && kind > ImageKind.Align)
                {
                    scanIndex += 1;
                    rawIndex = 0;
                    mIndex = 0;
                    index = idx;
                }

                switch (kind)
                {
                    case ImageKind.High:
                        await writer.WriteImage(path, kind, high).ConfigureAwait(false);
                        high += 1;
                        break;
                    case ImageKind.Align:
                        await writer.WriteImage(path, kind, align).ConfigureAwait(false);
                        align += 1;
                        break;
                    case ImageKind.Addr:
                        await writer.WriteImage(path, kind, scanIndex).ConfigureAwait(false);
                        break;
                    case ImageKind.Raw:
                        await writer.WriteRaw(path, scanIndex, rawIndex).ConfigureAwait(false);
                        rawIndex += 1;
                        break;
                    case ImageKind.MP:
                        await writer.WriteResult(path, v, scanIndex, mIndex).ConfigureAwait(false);
                        mIndex += 1;
                        break;
                }
            }
        }
        catch (Exception)
        {
            await writer.WriteFailure(scanIndex).ConfigureAwait(false);
            throw;
        }
        finally
        {
            await writer.WriteEndDate(DateTimeOffset.Now).ConfigureAwait(false);
        }
    }

    async Task<int> RunJobPlan(IJobContext ctx, JobPlan plan)
    {
        var options = await Align(ctx, plan.WaferMap, plan.Alignment).ConfigureAwait(false);
        var nav = stageService.WaferNavigator(options);
        return await RunImagingPipeline(ctx, nav, plan.ScanPoints, plan.Environment)
            .ConfigureAwait(false);
    }

    async Task<int> RunImagingPipeline(IJobContext ctx, IWaferNavigator nav, IReadOnlyList<ScanPoint> scanPoints, JobEnvironment env)
    {
        var pipeline = ImagingPipeline(ctx, nav, env);
        foreach (var p in scanPoints)
        {
            if (ctx.IsCanceled)
            {
                throw new OperationCanceledException();
            }

            try
            {
                var c = new Coordinate(p.X, p.Y, p.Dx, p.Dy, p.AddrOffsetX, p.AddrOffsetY);
                await pipeline.MoveTo(c, p.Pattern).ConfigureAwait(false);
            }
            catch (ThresholdException)
            {
                if (env.MatchFailureMode != MatchFailureMode.Ignore)
                {
                    throw;
                }
            }
        }
        return 0;
    }

    ImagingPipeline ImagingPipeline(IJobContext ctx, IWaferNavigator nav, JobEnvironment env)
    {
        return new ImagingPipeline(ctx, nav, semClient, imagingService)
        {
            MatchFailureMode = env.MatchFailureMode,
            IsDYNEnabled = env.IsDYNEnabled
        };
    }

    async Task<WaferNavigatorOptions> Align(IJobContext ctx, IWaferNavigator nav, Pattern pattern)
    {
        var pipeline = new ImagingPipeline(ctx, nav, semClient, imagingService)
        {
            MatchFailureMode = MatchFailureMode.Terminate,
            IsDYNEnabled = true,
        };
        var alignment = new WaferAlignment(ctx, imagingService, stageService, pipeline);
        var options = await alignment.Align(pattern).ConfigureAwait(false);
        await semClient.Settings.SetAngle(options.Angle);
        return options;
    }

    async Task<WaferNavigatorOptions> Align(IJobContext ctx, WaferMap map, List<Pattern> alignment)
    {
        await imagingService.UseHigh().ConfigureAwait(false);
        var options = await Align(ctx, stageService.WaferNavigator(map), alignment[0]).ConfigureAwait(false);
        await imagingService.UseSEM().ConfigureAwait(false);
        return await Align(ctx, stageService.WaferNavigator(options), alignment[1]).ConfigureAwait(false);
    }

    async Task<WaferNavigatorOptions> UnsafeAlign(IJobContext ctx, WaferMap map, Pattern high, Pattern align)
    {
        await imagingService.UnsafeUseHigh(false).ConfigureAwait(false);
        var options = await Align(ctx, stageService.WaferNavigator(map), high).ConfigureAwait(false);
        await imagingService.UnsafeUseSEM().ConfigureAwait(false);
        return await Align(ctx, stageService.WaferNavigator(options), align).ConfigureAwait(false);
    }

    static Pattern Copy(Pattern p)
    {
        return new Pattern
        {
            Kind = p.Kind,
            Ordinal = p.Ordinal,
            Dx = p.Dx,
            Dy = p.Dy,
            AddrOffsetX = p.AddrOffsetX,
            AddrOffsetY = p.AddrOffsetY,
            Reference = p.Reference,
            Addressing = p.Addressing,
            CDPoints = p.CDPoints,
            UseAddressing = p.UseAddressing,
            UseDifferentialDrive = p.UseDifferentialDrive
        };
    }
}